#include "hllc.hpp"
#include "sweFunctions.hpp"
#include "geometry.hpp"
#include "reconstruction.hpp"

namespace hms
{

template<typename MeshType>
void SWESolver<MeshType>::slopeSourceToFlux(
	Ref<Array3Ns<2>> flux,
	RowArray2s& zLR,
	const RowArray2s& dLR,
	const RowArray2s& hLR,
	const RowArray2s& zCellLR,
	const RowArray2s& dCellLR,
	const typename MeshBase<MeshType>::EdgeType& edge
) const {
	zLR = ( hLR < zLR ).select( hLR, zLR );

	// for (int i{0}; i<2; ++i){

	// 	flux.col(i).tail(2) += edge.normal().array() * phys::g / 2 *
	// 		(dLR[i] + dCellLR[i]) * (zLR[i] - zCellLR[i]);
	// }
	flux.template bottomRows<2>() +=
		edge.normal().array().template replicate<1,2>()
		* phys::g / 2
		* ( (dLR + dCellLR) * (zLR - zCellLR) ).template replicate<2,1>();
}


template<typename MeshType>
template<bool atomicWrite>
void SWESolver<MeshType>::fullyUpwind(const MeshBase<MeshType>& mesh, Index e){

	auto edge { mesh.edge(e) };
	Array2i lr { edge.cells() };

	if ( isDry()(0, lr).all() )
		return;
	
	auto getFieldVals = [&](const auto& field){
		return field.values()(hms::all, lr);
	};

	RowArray2s  dLR { getFieldVals ( d() ) };
	RowArray2s  zLR { getFieldVals ( z() ) };
	RowArray2s  hLR { zLR + dLR };
	Array2Ns<2> vLR { getFieldVals ( v() ) };

	Array2s vEdge { v().values()(hms::all, lr).rowwise().mean() };

	vEdge = transformCoords( vEdge, edge.normal() );

	Array3Ns<2> flux;
	if(vEdge[0] >= 0 ){
		flux.col(0)[0] = vEdge[0] * dLR[0];
		flux.col(0).template tail<2>() = vEdge[0] * dLR[0] * vLR.col(0);
	}
	else {
		flux.col(0)[0] = vEdge[0] * dLR[1];
		flux.col(0).template tail<2>() = vEdge[0] * dLR[1] * vLR.col(1);
	}

	flux.col(1) = flux.col(0);
	
	RowArray2s hNewLR;
	if( (dLR > minWaterDepth).all() || (dLR < minWaterDepth).all() ) 
		hNewLR = hLR.mean();
	else if(dLR[0] > minWaterDepth)
		hNewLR = 0.5 * (hLR[0] + std::min(hLR[0],zLR[1]));
	else
		hNewLR = 0.5 * (hLR[1] + std::min(hLR[1],zLR[0]));


	flux.row(1) += phys::g * dLR * hNewLR * edge.normal()[0];
	flux.row(2) += phys::g * dLR * hNewLR * edge.normal()[1];

	// Array3s flux;
	// flux[0] = dEdge * vEdge[0];
	// flux[1] = square(vEdge[0]) * dEdge + hllc::pressureFlux(dEdge);
	// flux[2] = vEdge.prod() * dEdge;
	registerFlux<atomicWrite>(flux.col(0), flux.col(1), edge, lr, mesh);
}


#define INSTANTIATE_FOR_ALL_MESHTYPES(MeshType) \
template class SWESolver<MeshType>; \
template void SWESolver<MeshType>::fullyUpwind<true >(const MeshBase<MeshType>&, Index); \
template void SWESolver<MeshType>::fullyUpwind<false>(const MeshBase<MeshType>&, Index);
#include "meshInstantiator.hpp"

} // namespace hms